/*
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#pragma once

#include "OgreMath/OgreVector3.h"
#include <vector>
#include <functional>
#include <ostream>
#include <iostream>
#include <unordered_map>
#include <unordered_set>
#include <boost/unordered_set.hpp>
#include <boost/unordered_map.hpp>
#include <list>
#include "NTree.h"
#include "Prerequisites.h"
#include "MeshSimplify.h"
#include "TriangleLookupTable.h"
#include "Vertex.h"

template<class VoxelData>
class MarchingCubes {
	static inline void getCubeVecs(Vector3i *vecs, const Vector3i &min)
	{
		vecs[0] = min;
		vecs[1] = min+Vector3i(0,0,1);
		vecs[2] = min+Vector3i(0,1,0);
		vecs[3] = min+Vector3i(0,1,1);
		vecs[4] = min+Vector3i(1,0,0);
		vecs[5] = min+Vector3i(1,0,1);
		vecs[6] = min+Vector3i(1,1,0);
		vecs[7] = min+Vector3i(1,1,1);
	}

	static int min2(int a, int b)
	{
		return (a < b) ? a : b;
	}

	typedef Vector3i CubeVecs[8];

	struct LeafCacheValue
	{
		LeafCacheValue() : children(0) {}
		LeafCacheValue(const LeafCacheValue& rhs)
		{
			operator=(rhs);
		}
		void operator=(const LeafCacheValue& rhs)
		{
			children = rhs.children;
			for (int i = 0; i < 8; i++) childrenData[i] = rhs.childrenData[i];
		}
		unsigned char children;
		VoxelData childrenData[8];
	};
	typedef std::vector<boost::unordered_map<Vector3i, LeafCacheValue, boost::Vector3iHasher> > LeafTables;
	typedef boost::unordered_set<Vector3i, boost::Vector3iHasher> LeafSet;

	__forceinline static void getNeighborCubes(LeafSet &minCubes,
		const Area<Vector3i> &area)
	{
		Box<Vector3i> box = area.toBox();
		for (int y = box.min.y-1; y <= box.max.y; y++)
		{
			for (int z = box.min.z-1; z <= box.max.z; z++)
			{
				minCubes.insert(Vector3i(box.min.x-1, y, z));
				minCubes.insert(Vector3i(box.max.x, y, z));
			}
		}
		for (int x = box.min.x-1; x <= box.max.x; x++)
		{
			for (int z = box.min.z-1; z <= box.max.z; z++)
			{
				minCubes.insert(Vector3i(x, box.min.y-1, z));
				minCubes.insert(Vector3i(x, box.max.y, z));
			}
		}
		for (int y = box.min.y-1; y <= box.max.y; y++)
		{
			for (int x = box.min.x-1; x <= box.max.x; x++)
			{
				minCubes.insert(Vector3i(x, y, box.min.z-1));
				minCubes.insert(Vector3i(x, y, box.max.z));
			}
		}
	}

	struct VertexIndexed
	{
		unsigned int index;
		VoxelData data;
	};
	static inline void processCube(std::unordered_map<Vector3i, VertexIndexed> &vertexMap,
		unsigned int &numVertices,
		std::vector<unsigned int> &indexBuffer,
		LeafTables &leafCache,
		const Vector3i &posMin)
	{
		CubeVecs cubeVecs;
		getCubeVecs(cubeVecs, posMin);
		unsigned char word = 0;
		VoxelData cubeCornerData[8];
		VoxelData* cubeData = nullptr;
		for (int i = 0; i < 8; i++)
		{
			bool isLeaf = queryIsLeaf(leafCache, Area<Vector3i>(0, cubeVecs[i]), cubeCornerData[i]);
			if (isLeaf)
			{
				if (!cubeData) cubeData = &cubeCornerData[i];
				word |= (1<<i);
			}
		}
		Vector3i offset = (posMin.doubleVec()+Vector3i(1,1,1));
		const std::vector<Triangle<Vector3i> >& triangles = TLT::getSingleton().table[word];
		for (unsigned int i = 0; i < triangles.size(); i++)
		{
			VertexIndexed indexed;
			indexed.data = *cubeData;	// for now one data for the whole cube, TODO: proper interpolation
			indexed.index = numVertices;
			Vector3i vertex1(triangles[i].p1+offset);
			auto tryInsert = vertexMap.insert(std::make_pair(vertex1, indexed));
			if (tryInsert.second) numVertices++;
			indexBuffer.push_back(tryInsert.first->second.index);

			indexed.index = numVertices;

			Vector3i vertex2(triangles[i].p2+offset);
			tryInsert = vertexMap.insert(std::make_pair(vertex2, indexed));
			if (tryInsert.second) numVertices++;
			indexBuffer.push_back(tryInsert.first->second.index);

			indexed.index = numVertices;

			Vector3i vertex3(triangles[i].p3+offset);
			tryInsert = vertexMap.insert(std::make_pair(vertex3, indexed));
			if (tryInsert.second) numVertices++;
			indexBuffer.push_back(tryInsert.first->second.index);
		}
	}

	__forceinline static bool isSolid(const typename NTree<Vector3i, VoxelData>::Node* node)
	{
		if (!node) return false;
		return node->getType() == Solid;
	}

	__forceinline static LeafCacheValue getNodeCacheValue(const typename NTree<Vector3i, VoxelData>::Node &node)
	{
		LeafCacheValue value;
		if (node.getType() == Nil) value.children = 0;
		if (node.getType() == Solid) value.children = 255;
		value.children = 0;
		for (int i = 0; i < 8; i++)
		{
			if (node.constGetSubNode(i).getType() == Solid) value.children |= (1<<i);
			value.childrenData[i] = node.constGetSubNode(i).getData();
		}
		return value;
	}
	__forceinline static unsigned char
	getRootCacheValue(NTree<Vector3i, VoxelData> &tree)
	{
		unsigned char cacheValue = 0;
		for (int i = 0; i < 8; i++)
		{
			if (tree.getRoot(i).getType() == Solid) cacheValue |= (1<<i);
		}
		return cacheValue;
	}

	// area query
	static bool queryIsLeaf(LeafTables &leafCache,
		Area<Vector3i> area,
		VoxelData& nodeData)
	{
		while (area.expo < leafCache.size()) {
			Vector3i parentPos(area.pos.half());
			auto cell = leafCache[area.expo].find(parentPos);
			if (cell != leafCache[area.expo].end())
			{
 				nodeData = cell->second.childrenData[area.getChildIndex()];
				return (cell->second.children & (1<<area.getChildIndex())) != 0;
			}
			area.expo++;
			area.pos = parentPos;
		}
		return false;
	}

	static bool queryIsLeaf(LeafTables &leafCache,
		Area<Vector3i> area)
	{
		while (area.expo < leafCache.size()) {
			Vector3i parentPos(area.pos.half());
			auto cell = leafCache[area.expo].find(parentPos);
			if (cell != leafCache[area.expo].end())
			{
				return (cell->second.children & (1<<area.getChildIndex())) != 0;
			}
			area.expo++;
			area.pos = parentPos;
		}
		return false;
	}

	// level 0 pos query
	/*static bool queryIsLeaf(LeafTables &leafCache,
		const Vector3i &pos,
		VoxelData& nodeData = VoxelData())
	{
		auto cell = leafCache[0].find(pos.half());
		if (cell != leafCache[0].end())
		{
			nodeData = cell->second.data[1<<Area<Vector3i>::getChildIndex(pos)];
			return (cell->second.children & (1<<Area<Vector3i>::getChildIndex(pos))) != 0;
		}
		return false;
	}*/

	/**
	Retrieves the set of cubes that must be triangulated for a node.
	*/
	static void nodeToCubes(
		LeafSet &minCubes,
		LeafTables &leafCache,
		typename NTree<Vector3i, VoxelData>::Node &node,
		const Area<Vector3i> &nodeArea,
		typename NTree<Vector3i, VoxelData>::Node solidDummyLeaf
		)
	{
		if (node.getType() == Both)
		{
			Area<Vector3i> subAreas[8];
			nodeArea.subAreas(subAreas);
			for (int i = 0; i < 8; i++)
				nodeToCubes(minCubes, leafCache, node.getSubNode(i), subAreas[i], solidDummyLeaf);
		}
		else if (node.getType() == Solid)
		{
			if (nodeArea.expo == 0)
			{
				// add cubes to minCubes list
				getNeighborCubes(minCubes, nodeArea);
			}
			else
			{
				unsigned int n = 0;
				Vector3i neighbors[Vector3i::neighborCount];
				nodeArea.pos.neighbors(neighbors);
				for (; n < Vector3i::neighborCount; n++)
				{
					if (!queryIsLeaf(leafCache, Area<Vector3i>(nodeArea.expo, neighbors[n])))
						break;
				}
				if (n == Vector3i::neighborCount) return;	// all neighbors are solid
				if (nodeArea.expo == 1)
				{
					LeafCacheValue value;
					value.children = 255;
					for (int i = 0; i < 8; i++) value.childrenData[i] = node.getData();
					leafCache[0].insert(std::make_pair(nodeArea.pos, value));	// add to cache
					//add surrounding cubes to minCubes
					getNeighborCubes(minCubes, nodeArea);
				}
				else
				{
					solidDummyLeaf.setData(node.getData());
					Area<Vector3i> subAreas[8];
					nodeArea.subAreas(subAreas);
					for (int i = 0; i < 8; i++)
						nodeToCubes(minCubes, leafCache, solidDummyLeaf, subAreas[i], solidDummyLeaf);
				}
			}
		}
	}

	static void buildLeafCache(
		LeafTables &leafCache,
		typename NTree<Vector3i, VoxelData>::Node &node,
		const Area<Vector3i> &nodeArea)
	{
		if (node.getType() == Both)
		{
			vAssert(nodeArea.expo > 0);
			LeafCacheValue cacheValue = getNodeCacheValue(node);
			if (cacheValue.children != 0) leafCache[nodeArea.expo-1].insert(std::make_pair(nodeArea.pos, cacheValue));
			Area<Vector3i> subAreas[8];
			nodeArea.subAreas(subAreas);
			for (int i = 0; i < 8; i++)
				buildLeafCache(leafCache, node.getSubNode(i), subAreas[i]);
		}
	}

	static void buildLeafCache(
		LeafTables &leafCache,
		NTree<Vector3i, VoxelData> &tree)
	{
		Vector3i neighbors[8];
		Vector3i pos(-1);
		pos.positiveNeighbors(neighbors);
		for (int i = 0; i < 8; i++) {
			Area<Vector3i> subArea(tree.getSize(), neighbors[i]);
			buildLeafCache(leafCache, tree.getRoot(i), subArea);
		}
	}

public:

	template<class T>
	static void computeVertexNormals(std::vector<Ogre::Vector3> &normalBuffer, const std::vector<Vertex<T> > &vertexBuffer, const std::vector<unsigned int> &indexBuffer)
	{
		normalBuffer.resize(vertexBuffer.size());
		std::vector<std::vector<Ogre::Vector3>> vertexNormalNeighbors;
		vertexNormalNeighbors.resize(vertexBuffer.size());
		for (int i = 0; i < ((int)indexBuffer.size())-3; i+=3)
		{
			Ogre::Vector3 normal((vertexBuffer[indexBuffer[i+2]].position-vertexBuffer[indexBuffer[i]].position).
				crossProduct((vertexBuffer[indexBuffer[i+1]].position-vertexBuffer[indexBuffer[i]].position)));
			normal.normalise();
			vertexNormalNeighbors[indexBuffer[i]].push_back(normal);
			vertexNormalNeighbors[indexBuffer[i+1]].push_back(normal);
			vertexNormalNeighbors[indexBuffer[i+2]].push_back(normal);
		}
		for (unsigned int i = 0; i < normalBuffer.size(); i++)
		{
			normalBuffer[i] = Ogre::Vector3(0,0,0);
			for (unsigned int n = 0; n < vertexNormalNeighbors[i].size(); n++)
				normalBuffer[i] += vertexNormalNeighbors[i][n];
			normalBuffer[i].normalise();
		}
	}

	template<class T>
	static void smoothMesh(std::vector<Vertex<T> > &vertexBuffer, const std::vector<unsigned int> &indexBuffer, unsigned int numIterations)
	{
		std::vector<std::unordered_set<unsigned int>> vertexNeighbors;		
		vertexNeighbors.resize(vertexBuffer.size());
		for (int i = 0; i < ((int)indexBuffer.size())-2; i+=3)
		{
			vertexNeighbors[indexBuffer[i]].insert(indexBuffer[i+1]);
			vertexNeighbors[indexBuffer[i]].insert(indexBuffer[i+2]);

			vertexNeighbors[indexBuffer[i+1]].insert(indexBuffer[i]);
			vertexNeighbors[indexBuffer[i+1]].insert(indexBuffer[i+2]);

			vertexNeighbors[indexBuffer[i+2]].insert(indexBuffer[i]);
			vertexNeighbors[indexBuffer[i+2]].insert(indexBuffer[i+1]);
		}

		std::vector<Vertex<T> > swapBuffer = vertexBuffer;
		std::vector<Vertex<T> > *oldBufferPtr = &vertexBuffer;
		std::vector<Vertex<T> > *newBufferPtr = &swapBuffer;
		for (unsigned int i = 0; i < numIterations; i++)
		{
			for (unsigned int v = 0; v < oldBufferPtr->size(); v++)
			{
				Ogre::Vector3 accum = (*oldBufferPtr)[v].position;
				for (auto it = vertexNeighbors[v].begin(); it != vertexNeighbors[v].end(); it++)
					accum += (*oldBufferPtr)[*it].position;
				(*newBufferPtr)[v].position = accum / (1.0f+vertexNeighbors[v].size());
			}
			std::swap(oldBufferPtr, newBufferPtr);
		}
		if (oldBufferPtr != &vertexBuffer) vertexBuffer = swapBuffer;
	}

	template<class MeshWriter>
	static void treeToMesh(const std::string &fileName, NTree<Vector3i, VoxelData> &tree,
			float voxelsPerUnit, int numSmoothIterations = 1, bool bSimplify=true)
	{
		LeafSet minCubes;
		std::cout << "[Marching cubes] Tree size is " << (int)tree.getSize() << std::endl;
		unsigned int maxBuckets = (1<(tree.getSize()+1))*(1<<(tree.getSize()+1))*10;
		if (maxBuckets > 10000000) maxBuckets = 10000000;
		minCubes.rehash(maxBuckets);
		LeafTables leafCache;
		leafCache.resize(tree.getSize());
		for (unsigned int i = 0; i < leafCache.size(); i++) {
			int dim = 1 << (leafCache.size() + 1 - i);
			unsigned int numBuckets = dim*dim*10;
			if (numBuckets > maxBuckets) numBuckets = maxBuckets;
			leafCache[i].rehash(numBuckets);
		}

		std::cout << "[Marching cubes] Building leaf cache..." << std::endl;
		buildLeafCache(leafCache, tree);

		/*for (auto i = leafCache.begin(); i != leafCache.end(); i++)
		{
			int maxBucket = 0;
			for (int b = 0; b < i->bucket_count(); b++)
				if (i->bucket_size(b) > maxBucket) maxBucket = i->bucket_size(b);
			std::cout << maxBucket << std::endl;
		}*/

		std::cout << "[Marching cubes] Fetching cubes..." << std::endl;
		Vector3i neighbors[8];
		Vector3i pos(-1);
		pos.positiveNeighbors(neighbors);
		for (int i = 0; i < 8; i++) {
			Area<Vector3i> subArea(tree.getSize(), neighbors[i]);
			nodeToCubes(minCubes, leafCache, tree.getRoot(i), subArea, typename NTree<Vector3i, VoxelData>::Node(Solid));
		}

		std::cout << "[Marching cubes] Processing fetched cubes (" << minCubes.size() << ") ..." << std::endl;

		std::unordered_map<Vector3i, VertexIndexed> vertexMap;
		std::vector<unsigned int> indexBuffer;
		unsigned int numVertices = 0;

		for (auto it = minCubes.begin(); it != minCubes.end(); it++)
			processCube(vertexMap, numVertices, indexBuffer, leafCache, *it);

		std::cout << "[Marching cubes] " << vertexMap.size() << " vertices created." << std::endl;

		// each marching cubes cube has size 2
		std::vector<Vertex<VoxelData> > vertexBuffer;
		vertexBuffer.resize(vertexMap.size());
		for (auto it = vertexMap.begin(); it != vertexMap.end(); it++)
		{
			vertexBuffer[it->second.index].position = it->first.toOgreVec()*0.5f;
			vertexBuffer[it->second.index].data = it->second.data;
		}

		if(bSimplify)
			MeshSimplify::simplify(vertexBuffer, indexBuffer);

		if (numSmoothIterations > 0)
		{
			std::cout << "[Marching cubes] Smoothing mesh ..." << std::endl;
			smoothMesh(vertexBuffer, indexBuffer, numSmoothIterations);
		}

		//DEBUG
		/*std::vector<std::pair<Vector3i, VoxelData>> lPoints = NTreeHelpers::getPoints(tree);
		for(auto i=lPoints.begin(); i!=lPoints.end(); i++)
			printf("Material :%d\n", (unsigned int) i->second);*/

		// finally scale with respect to voxelsPerUnit 
		float scale = 1.0f / voxelsPerUnit;
		for (auto it = vertexBuffer.begin(); it != vertexBuffer.end(); it++)
			it->position *= scale;

		std::cout << "[Marching cubes] Computing normals ..." << std::endl;
		std::vector<Ogre::Vector3> normalBuffer;
		computeVertexNormals(normalBuffer, vertexBuffer, indexBuffer);

		std::cout << "[Marching cubes] Writing file \"" << fileName << "\" ..." << std::endl;
		MeshWriter::writeMesh(fileName, vertexBuffer, normalBuffer, indexBuffer);
		std::cout << "[Marching cubes] Finished!" << std::endl;
	}
};
